package cn.iocoder.yudao.module.aiposter.remote;

import cn.iocoder.yudao.framework.common.util.json.JsonUtils;
import cn.iocoder.yudao.module.aiposter.config.RemoteConfig;
import cn.iocoder.yudao.module.aiposter.remote.vo.*;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.collections4.map.HashedMap;
import org.springframework.core.ParameterizedTypeReference;
import org.springframework.http.*;
import org.springframework.stereotype.Component;
import org.springframework.web.client.RestTemplate;

import javax.annotation.Resource;
import java.util.ArrayList;
import java.util.List;
import java.util.Map;

@Component
@Slf4j
public class VecClient {

    @Resource
    private RemoteConfig remoteConfig;
    @Resource
    private RestTemplate restTemplate;

    private final String PATH_SIMILARITY_SEARCH = "/poster_similarity_search";
    private final String PATH_VEC_ADD = "/poster_insert";
    private final String PATH_VEC_CLEAN_ALL = "/poster_clean";
    private final String PATH_VEC_UPDATE = "/poster_update";
    private final String PATH_VEC_DELETE = "/poster_delete";

    /**
     * 请求向量库获取列表中的关联数据字段
     *
     * @param requestItem 待匹配字段
     * @return 相似匹配内容
     */
    public List<VecDBResp.SimilaritySearchItem> similaritySearch(List<String> requestItem, Integer numResult, PosterFilterBO posterFilterBO) {
        try {
            Map<String, Object> request = new HashedMap<>();
            request.put("query_text_list", requestItem);
            if (null != posterFilterBO) {
                request.put("query_filter", posterFilterBO);
            }
            request.put("n_results", numResult);
            ParameterizedTypeReference reference = new ParameterizedTypeReference<VecDBBaseResp<List<VecDBResp.SimilaritySearchItem>>>() {};
            return vecCall(PATH_SIMILARITY_SEARCH, request, HttpMethod.POST, reference);
        } catch (Exception e) {
            e.printStackTrace();
        }
        return null;
    }

    /**
     * 向量库添加数据
     *
     * @param requestItem 添加列表
     * @return 是否添加成功
     */
    public boolean add(VecDBItemReq requestItem) {
        return add(new ArrayList<VecDBItemReq>() {{
            add(requestItem);
        }});
    }

    /**
     * 向量库添加数据
     *
     * @param requestItem 添加列表
     * @return 是否添加成功
     */
    public boolean add(List<VecDBItemReq> requestItem) {
        try {
            vecCall(PATH_VEC_ADD, requestItem, HttpMethod.POST);
            return true;
        } catch (Exception e) {
            e.printStackTrace();
        }
        return false;
    }

    /**
     * 向量库删除数据
     *
     * @return 是否删除成功
     */
    public boolean cleanAll() {
        try {
            vecCall(PATH_VEC_CLEAN_ALL, null, HttpMethod.POST);
            return true;
        } catch (Exception e) {
            e.printStackTrace();
        }
        return false;
    }

    /**
     * 向量库更新数据
     *
     * @param requestItem 更新列表
     * @return 是否删除成功
     */
    public boolean update(VecDBItemReq requestItem) {
        try {
            vecCall(PATH_VEC_UPDATE, requestItem, HttpMethod.POST);
            return true;
        } catch (Exception e) {
            e.printStackTrace();
        }
        return false;
    }


    /**
     * 向量库删除某个数据
     *
     * @param requestItem 更新列表
     * @return 是否删除成功
     */
    public boolean delete(VecDBItemReq requestItem) {
        try {
            vecCall(PATH_VEC_DELETE, requestItem, HttpMethod.POST);
            return true;
        } catch (Exception e) {
            e.printStackTrace();
        }
        return false;
    }

    private <T> T vecCall(String path, Object requestObject, HttpMethod method) {
        return vecCall(path, requestObject, method, null);
    }

    /**
     * 访问向量库
     *
     * @param path
     * @param requestObject
     * @param method
     * @param reference     如果返回类型是List类型需要自己指定reference
     * @param <T>
     * @return
     */
    private <T> T vecCall(String path, Object requestObject, HttpMethod method, ParameterizedTypeReference reference) {
        String questUrl = remoteConfig.getVecHost() + path;
        long startTime = System.currentTimeMillis();
        if (null != requestObject) {
            log.debug("requestInfo :{}", requestObject.toString());
        }
        try {
            HttpHeaders headers = new HttpHeaders();
            headers.setContentType(MediaType.APPLICATION_JSON);
            HttpEntity<String> entity = new HttpEntity<>(JsonUtils.toJsonString(requestObject), headers);
            if (null == reference) {
                reference = new ParameterizedTypeReference<VecDBBaseResp<T>>() {
                };
            }
            ResponseEntity<VecDBBaseResp<T>> responseEntity = restTemplate.exchange(questUrl, method, entity,
                    reference);

            log.debug("VecClient response,cost:{}, questUrl:{} \n entity:{} \n body:{}",
                    System.currentTimeMillis() - startTime, questUrl, entity, responseEntity.getBody());
            VecDBBaseResp<T> resp = responseEntity.getBody();
            if (null != resp && null != resp.getCode() && resp.getCode().equals(200)) {
                return resp.getData();
            } else {
                throw new RuntimeException("vec call result error {}" + resp);
            }
        } catch (Exception e) {
            log.error("requestCwVecDB error", e);
        }
        return null;
    }
}
